# -*- coding: utf-8 -*-
"""
LLM 提示词模板生成器 - 动态生成基于模型类型和参数模式的提示词

@author: AI Assistant
"""
from typing import Dict, Any, List
from loguru import logger


class PromptTemplateGenerator:
    """提示词模板生成器"""
    
    # 模型类型中文名称映射
    MODEL_TYPE_NAMES = {
        "dhf": "大伙房模型",
        "xaj": "新安江模型",
        "tank": "Tank模型",
        # 可以添加更多模型类型
    }
    
    def __init__(self):
        """初始化提示词生成器"""
        pass
    
    def generate_system_prompt(
        self,
        model_type: str,
        stcd: str,
        stcd_name: str,
        para_schema: Dict[str, Any],
        initial_state: Dict[str, Any],
        parameters: Dict[str, Any]
    ) -> str:
        """根据模型信息动态生成系统提示词
        
        Args:
            model_type: 模型类型
            stcd: 站点代码
            stcd_name: 站点名称
            para_schema: 参数模式定义
            initial_state: 当前初始状态
            parameters: 当前模型参数
            
        Returns:
            完整的系统提示词
        """
        # 获取模型类型名称
        model_name = self.MODEL_TYPE_NAMES.get(model_type, model_type.upper())
        
        # 生成参数信息部分
        param_info = self._generate_parameter_info(para_schema, initial_state, parameters)
        
        # 构建系统提示词
        system_prompt = f"""你是一个专业的水文模型参数调整助手。你的任务是帮助用户通过自然语言对话来修改水文模型的参数。

## 当前模型信息

- **模型类型**: {model_name} ({model_type})
- **站点代码**: {stcd}
- **站点名称**: {stcd_name}

## 可调整的参数

{param_info}

## 你的职责

1. **理解用户意图**: 仔细分析用户的自然语言描述，理解他们想要修改哪些参数
2. **参数转换**: 将用户的描述转换为具体的参数路径和数值
3. **如实返回**: 按用户要求返回参数值，系统会自动过滤超出范围的参数
4. **清晰解释**: 向用户解释你所做的参数修改

## 响应格式要求

**重要**: 请直接输出 JSON，不要输出思考过程、分析步骤或任何额外说明。你必须**严格**使用以下 JSON 格式回复：

```json
{{
    "understood": true/false,  // 是否理解用户意图
    "message": "你对用户的回复，用中文",  // 友好的回复消息
    "parameter_changes": {{  // 参数修改建议（如果理解了用户意图）
        "initial_state.parameter_name": 新值,
        "parameters.parameter_name": 新值
    }},
    "explanation": "详细解释你为什么这样修改参数"
}}
```

## 示例对话

**示例 1 - 正常范围**:
**用户**: "请将初始水位提高 10cm"
**你的回复**:
```json
{{
    "understood": true,
    "message": "好的，我将初始水位从 100.00 米提高到 100.10 米（增加了 10 厘米）",
    "parameter_changes": {{
        "initial_state.water_level": 100.10
    }},
    "explanation": "根据您的要求，我将初始水位参数从原来的 100.00 米增加了 0.10 米（10 厘米）。"
}}
```

**示例 2 - 用户请求的值超出范围（按用户要求如实返回，系统会自动处理）**:
**用户**: "请将 A 参数设置为 300"
**你的回复**:
```json
{{
    "understood": true,
    "message": "好的，我将 A 参数设置为 300",
    "parameter_changes": {{
        "parameters.A": 300
    }},
    "explanation": "已按您的要求将 A 参数设置为 300"
}}
```
**注意**: 即使 300 可能超出范围，也按用户要求返回。系统会自动过滤超出范围的参数。

**示例 3 - 多个参数，部分可能超出范围**:
**用户**: "请将 A 设为 300，B 设为 0.5"
**你的回复**:
```json
{{
    "understood": true,
    "message": "好的，我将 A 参数设置为 300，B 参数设置为 0.5",
    "parameter_changes": {{
        "parameters.A": 300,
        "parameters.B": 0.5
    }},
    "explanation": "已按您的要求修改了 A 和 B 两个参数"
}}
```
**注意**: 全部返回，系统会自动过滤掉超出范围的参数。

## 注意事项

- **直接输出 JSON，不要输出任何思考过程或额外说明**
- 始终以 JSON 格式回复，不要添加额外的文字说明
- 如果用户的请求不清楚，设置 `understood: false` 并在 `message` 中询问具体信息
- 参数路径格式必须是 `initial_state.参数名` 或 `parameters.参数名`
- **按用户要求如实返回参数值**，无需担心范围问题，系统会自动处理超出范围的参数
- 使用中文与用户交流，保持友好和专业
"""
        
        return system_prompt
    
    def _generate_parameter_info(
        self,
        para_schema: Dict[str, Any],
        initial_state: Dict[str, Any],
        parameters: Dict[str, Any]
    ) -> str:
        """生成参数信息文本
        
        Args:
            para_schema: 参数模式定义
            initial_state: 当前初始状态
            parameters: 当前模型参数
            
        Returns:
            格式化的参数信息文本
        """
        param_lines = []
        
        # 处理初始状态参数
        if "initial_state" in para_schema:
            param_lines.append("### 初始状态参数 (initial_state)\n")
            for param_name, param_def in para_schema["initial_state"].items():
                current_value = initial_state.get(param_name, "未设置")
                param_desc = self._format_parameter_description(
                    param_name, param_def, current_value, "initial_state"
                )
                param_lines.append(param_desc)
        
        # 处理模型参数
        if "parameters" in para_schema:
            param_lines.append("\n### 模型参数 (parameters)\n")
            for param_name, param_def in para_schema["parameters"].items():
                current_value = parameters.get(param_name, "未设置")
                param_desc = self._format_parameter_description(
                    param_name, param_def, current_value, "parameters"
                )
                param_lines.append(param_desc)
        
        if not param_lines:
            return "（无参数定义）"
        
        return "\n".join(param_lines)
    
    def _format_parameter_description(
        self,
        param_name: str,
        param_def: Dict[str, Any],
        current_value: Any,
        category: str
    ) -> str:
        """格式化单个参数的描述
        
        Args:
            param_name: 参数名称
            param_def: 参数定义
            current_value: 当前值
            category: 参数类别（initial_state 或 parameters）
            
        Returns:
            格式化的参数描述
        """
        desc = param_def.get("description", "无描述")
        param_type = param_def.get("type", "unknown")
        unit = param_def.get("unit", "")
        min_val = param_def.get("min")
        max_val = param_def.get("max")
        
        # 构建参数描述
        parts = [
            f"- **{param_name}** ({category}.{param_name})",
            f"  - 描述: {desc}",
            f"  - 类型: {param_type}",
            f"  - 当前值: {current_value} {unit}".strip()
        ]
        
        if min_val is not None or max_val is not None:
            range_str = f"  - 取值范围: "
            if min_val is not None and max_val is not None:
                range_str += f"[{min_val}, {max_val}] {unit}".strip()
            elif min_val is not None:
                range_str += f">= {min_val} {unit}".strip()
            elif max_val is not None:
                range_str += f"<= {max_val} {unit}".strip()
            parts.append(range_str)
        
        return "\n".join(parts)
    
    def generate_user_prompt(self, user_message: str, conversation_history: List[Dict[str, str]] = None) -> str:
        """生成用户提示词（包含对话历史）
        
        Args:
            user_message: 用户当前消息
            conversation_history: 对话历史列表，格式：[{"role": "user/assistant", "content": "..."}]
            
        Returns:
            完整的用户提示词
        """
        if not conversation_history:
            return user_message
        
        # 构建包含历史的提示词
        prompt_parts = ["## 对话历史\n"]
        for msg in conversation_history[-5:]:  # 只保留最近5轮对话
            role = "用户" if msg["role"] == "user" else "助手"
            prompt_parts.append(f"**{role}**: {msg['content']}\n")
        
        prompt_parts.append(f"\n## 当前用户消息\n\n{user_message}")
        
        return "\n".join(prompt_parts)


# 使用示例
if __name__ == "__main__":
    import sys
    from pathlib import Path
    
    # 添加项目根目录到路径
    project_root = Path(__file__).resolve().parent.parent
    if str(project_root) not in sys.path:
        sys.path.insert(0, str(project_root))
    
    from services.session_service import SessionService
    
    generator = PromptTemplateGenerator()
    session_service = SessionService()
    
    # 测试会话 ID（可以从命令行参数获取）
    test_session_id = "8c73b607-8fd4-4d7c-ae45-4b65a78b287b"
    if len(sys.argv) > 1:
        test_session_id = sys.argv[1]
    
    print(f"正在加载会话: {test_session_id}")
    print("=" * 80)
    
    # 从数据库获取会话数据
    session = session_service.get_session_by_id(test_session_id)
    
    if session:
        print(f"✅ 成功加载会话:")
        print(f"  - 模型类型: {session.model_type}")
        print(f"  - 站点: {session.stcd} ({session.stcd_name})")
        print(f"  - 参数数量: {len(session.parameters)}")
        print(f"  - 初始状态参数数量: {len(session.initial_state)}")
        print("=" * 80)
        
        # 生成系统提示词
        system_prompt = generator.generate_system_prompt(
            model_type=session.model_type,
            stcd=session.stcd,
            stcd_name=session.stcd_name or "",
            para_schema=session.para_schema,
            initial_state=session.initial_state,
            parameters=session.parameters
        )
        
        print("\n生成的系统提示词:")
        print("=" * 80)
        print(system_prompt)
        print("=" * 80)
        
        # 显示参数信息
        print("\n参数详情:")
        print("-" * 80)
        if session.initial_state:
            print("\n初始状态参数:")
            for key, value in list(session.initial_state.items())[:5]:
                print(f"  - {key}: {value}")
            if len(session.initial_state) > 5:
                print(f"  ... 还有 {len(session.initial_state) - 5} 个参数")
        
        if session.parameters:
            print("\n模型参数:")
            for key, value in list(session.parameters.items())[:5]:
                print(f"  - {key}: {value}")
            if len(session.parameters) > 5:
                print(f"  ... 还有 {len(session.parameters) - 5} 个参数")
    else:
        print(f"❌ 未找到会话: {test_session_id}")
        print("\n提示: 使用方法: python prompt_templates.py [session_id]")

